# frozen_string_literal: true

module Gitlab
  module Graphql
    module Aggregations
      module VulnerabilityStatistics
        class LazyAggregate
          include ::Gitlab::Graphql::Deferred

          attr_reader :vulnerable, :lazy_state

          def initialize(query_ctx, vulnerable, filter: nil, include_subgroups: false)
            @vulnerable = vulnerable.respond_to?(:sync) ? vulnerable.sync : vulnerable
            @include_subgroups = include_subgroups
            @filter = filter
            # Initialize the loading state for this query,
            # or get the previously-initiated state
            @lazy_state = query_ctx[:lazy_vulnerability_statistics_aggregate] ||= {
              pending_vulnerables: Hash.new { |h, k| h[k] = Set.new },
              loaded_objects: Hash.new { |h, k| h[k] = {} }
            }
            # Register this ID to be loaded later:
            @lazy_state[:pending_vulnerables][key] << vulnerable
          end

          # Return the loaded record, hitting the database if needed
          def execute
            # Check if the records are not loaded yet
            if @lazy_state[:pending_vulnerables][key].present?
              load_records_into_loaded_objects
            end

            @lazy_state[:loaded_objects][key][@vulnerable]
          end

          private

          def key
            [@include_subgroups, @filter]
          end

          def load_records_into_loaded_objects
            # The record hasn't been loaded yet, so
            # hit the database with all pending IDs to prevent N+1
            @lazy_state[:pending_vulnerables].each do |(include_subgroups, filter), pending_vulnerables|
              next if pending_vulnerables.blank?

              key = [include_subgroups, filter]
              grades = ::Vulnerabilities::ProjectsGrade.grades_for(pending_vulnerables, filter: filter, include_subgroups: include_subgroups)

              pending_vulnerables.each do |vulnerable|
                @lazy_state[:loaded_objects][key][vulnerable] = grades[vulnerable]
              end

              @lazy_state[:pending_vulnerables][key].clear
            end
          end
        end
      end
    end
  end
end
